import os
import sys
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Rectangle, Circle, Polygon, FancyArrowPatch
import glob

# 解决OpenMP警告（如果存在）
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# === 设置论文级别的绘图样式 ===
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times New Roman', 'DejaVu Serif']
plt.rcParams['font.size'] = 11
plt.rcParams['axes.linewidth'] = 1.5
plt.rcParams['axes.labelsize'] = 12
plt.rcParams['axes.titlesize'] = 13
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10
plt.rcParams['legend.fontsize'] = 10
plt.rcParams['legend.framealpha'] = 0.95
plt.rcParams['legend.edgecolor'] = 'gray'
plt.rcParams['grid.alpha'] = 0.3
plt.rcParams['grid.linestyle'] = '--'
plt.rcParams['lines.linewidth'] = 2.0

# === 1. 通过模型目录读取所有CSV文件 ===
# 用法: python render.py <rollout_directory>
# 例如: python render.py runs/a2c_continuous_moose_default_1124_2055/rollouts/20251125-143000
# 或者: python render.py runs/a2c_continuous_moose_default_1124_2055  (自动查找最新的rollouts子文件夹)

if len(sys.argv) > 1:
    input_dir = sys.argv[1]
    print(f"Using input directory: {input_dir}")
else:
    # 默认目录
    input_dir = "runs/a2c_continuous_moose_default_1126_2104"
    print(f"No directory specified, using default: {input_dir}")

# 智能查找CSV文件：
# 1. 如果输入目录直接包含CSV文件，使用该目录
# 2. 如果输入目录包含rollouts子文件夹，自动查找最新的rollout子文件夹
csv_files = glob.glob(os.path.join(input_dir, "*.csv"))
model_dir = input_dir

if not csv_files:
    # 尝试在rollouts子文件夹中查找
    rollouts_dir = os.path.join(input_dir, "rollouts")
    if os.path.exists(rollouts_dir):
        # 找到所有rollout子文件夹，按修改时间排序
        rollout_subdirs = [d for d in glob.glob(os.path.join(rollouts_dir, "*")) if os.path.isdir(d)]
        if rollout_subdirs:
            # 使用最新的rollout文件夹
            latest_rollout = max(rollout_subdirs, key=os.path.getmtime)
            csv_files = glob.glob(os.path.join(latest_rollout, "*.csv"))
            if csv_files:
                model_dir = latest_rollout
                print(f"Found rollouts subdirectory, using latest: {model_dir}")
            else:
                raise FileNotFoundError(f"No CSV files found in latest rollout directory: {latest_rollout}")
        else:
            raise FileNotFoundError(f"No rollout subdirectories found in: {rollouts_dir}")
    else:
        raise FileNotFoundError(f"No CSV files found in directory: {input_dir}")

print(f"Found {len(csv_files)} CSV files in {model_dir}")

# 列名定义
col_names_32 = ["x", "y", "psi", "x_dot", "y_dot", "psi_dot", "delta",
                "omega_fr", "omega_fl", "omega_rr", "omega_rl", "f1", "f2", "f3", "f4", 
                "r", "beta", "v", "Zs", "phi", "theta", 
                "Z11", "Z12", "Z13", "Z14", 
                "dZs", "dphi", "dtheta", "dZ11", "dZ12", "dZ13", "dZ14"]
col_names_extra = ["LLTR_front", "LLTR_rear", "gamma_fr", "gamma_fl", "gamma_rr", "gamma_rl"]
col_names_38 = col_names_32 + col_names_extra

# 加载所有CSV数据
all_data = []
for csv_file in csv_files:
    try:
        df = pd.read_csv(csv_file, header=None)
        num_cols = len(df.columns)

        if num_cols >= len(col_names_38):
            df.columns = col_names_38[:num_cols] + [f"col_{i}" for i in range(len(col_names_38), num_cols)]
        elif num_cols == 32:
            df.columns = col_names_32
        elif num_cols >= 19:
            df.columns = col_names_32[:num_cols] + [f"col_{i}" for i in range(len(col_names_32), num_cols)]
        else:
            print(f"Skipping {csv_file}: insufficient columns ({num_cols})")
            continue
        
        all_data.append({
            'file': csv_file,
            'df': df,
            'x': df["x"].values,
            'y': df["y"].values,
            'phi': np.degrees(df["phi"].values),  # 转换为角度
            'theta': np.degrees(df["theta"].values),  # 转换为角度
            'beta': np.degrees(df["beta"].values) if "beta" in df.columns else np.zeros(len(df)),
        })
        print(f"  Loaded: {os.path.basename(csv_file)} ({len(df)} steps)")
    except Exception as e:
        print(f"  Error loading {csv_file}: {e}")

if not all_data:
    raise ValueError("No valid CSV files could be loaded")

print(f"\nSuccessfully loaded {len(all_data)} CSV files")

# === 2. 计算统计量 ===
# 找到最大长度用于对齐
max_len = max(len(d['x']) for d in all_data)
print(f"Max trajectory length: {max_len} steps")

# 将所有数据对齐到相同长度（用NaN填充短的轨迹）
def align_data(data_list, key, max_len):
    """将不同长度的数据对齐，短的用NaN填充"""
    aligned = np.full((len(data_list), max_len), np.nan)
    for i, d in enumerate(data_list):
        length = len(d[key])
        aligned[i, :length] = d[key]
    return aligned

x_aligned = align_data(all_data, 'x', max_len)
phi_aligned = align_data(all_data, 'phi', max_len)
theta_aligned = align_data(all_data, 'theta', max_len)
beta_aligned = align_data(all_data, 'beta', max_len)

# 计算均值和标准差（忽略NaN）
x_mean = np.nanmean(x_aligned, axis=0)
x_std = np.nanstd(x_aligned, axis=0)
phi_mean = np.nanmean(phi_aligned, axis=0)
phi_std = np.nanstd(phi_aligned, axis=0)
theta_mean = np.nanmean(theta_aligned, axis=0)
theta_std = np.nanstd(theta_aligned, axis=0)
beta_mean = np.nanmean(beta_aligned, axis=0)
beta_std = np.nanstd(beta_aligned, axis=0)

# 时间步
frames = np.arange(max_len)

# === 3. 创建图形 ===
fig = plt.figure(figsize=(14, 12), dpi=120)
gs = fig.add_gridspec(4, 1, height_ratios=[1, 1, 1, 1], hspace=0.35)

# 为所有子图设置背景色
subplot_bg = '#fafafa'

# 定义专业配色
color_individual = '#3498db'  # 蓝色，用于单条曲线
color_mean = '#e74c3c'        # 红色，用于均值曲线
color_ci = '#e74c3c'          # 红色，用于置信区间

# === 子图1: 存活距离图（水平条形图）===
# 每条轨迹显示为一条水平线，线的长度代表该轨迹最终到达的x位置（存活距离）
ax_x = fig.add_subplot(gs[0])
ax_x.set_facecolor(subplot_bg)
ax_x.set_xlabel("Final X Position (m)", fontsize=11, fontweight='bold')
ax_x.set_ylabel("Trajectory Index", fontsize=11, fontweight='bold')
ax_x.set_title(f"(a) Survival Distance - {len(all_data)} Trajectories", fontsize=12, fontweight='bold', pad=10)
ax_x.grid(True, alpha=0.35, linewidth=0.8, which='major', axis='x')
ax_x.minorticks_on()
ax_x.grid(True, alpha=0.15, linewidth=0.5, which='minor', axis='x')

# 获取每条轨迹的最终x位置（存活距离）
final_x_positions = [d['x'][-1] for d in all_data]

# 按存活距离排序（从小到大），方便观察
sorted_indices = np.argsort(final_x_positions)
sorted_final_x = [final_x_positions[i] for i in sorted_indices]

# 绘制水平条形图：每条轨迹一条水平线
num_traj = len(all_data)
y_positions = np.arange(num_traj)

# 使用颜色映射表示存活距离
cmap = plt.cm.RdYlGn  # 红-黄-绿色映射（短=红，长=绿）
max_x = max(final_x_positions)
min_x = min(final_x_positions)
norm = plt.Normalize(vmin=min_x, vmax=max_x)

for i, (y_pos, x_final) in enumerate(zip(y_positions, sorted_final_x)):
    color = cmap(norm(x_final))
    ax_x.barh(y_pos, x_final, height=0.7, color=color, edgecolor='none', alpha=0.85)

# 添加均值线
mean_final_x = np.mean(final_x_positions)
ax_x.axvline(x=mean_final_x, color=color_mean, linewidth=2.5, linestyle='--', 
             label=f'Mean: {mean_final_x:.1f} m')

# 设置坐标轴范围
ax_x.set_xlim(0, 100)
ax_x.set_ylim(-0.5, num_traj - 0.5)
ax_x.set_yticks(y_positions[::max(1, num_traj//10)])  # 每隔几个显示一个刻度
ax_x.legend(loc='lower right', fontsize=10, framealpha=0.95)

# 添加统计信息文本
stats_text = f'Min: {min_x:.1f}m | Mean: {mean_final_x:.1f}m | Max: {max_x:.1f}m'
ax_x.text(0.02, 0.95, stats_text, transform=ax_x.transAxes, fontsize=9,
          verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.9))

# === 子图2: Roll Angle (phi) ===
ax_phi = fig.add_subplot(gs[1])
ax_phi.set_facecolor(subplot_bg)
ax_phi.set_xlabel("Time Step", fontsize=11, fontweight='bold')
ax_phi.set_ylabel("Roll Angle, $\\phi$ (deg)", fontsize=11, fontweight='bold')
ax_phi.set_title(f"(b) Vehicle Roll Angle - {len(all_data)} Trajectories", fontsize=12, fontweight='bold', pad=10)
ax_phi.grid(True, alpha=0.35, linewidth=0.8, which='major')
ax_phi.minorticks_on()
ax_phi.grid(True, alpha=0.15, linewidth=0.5, which='minor')

# 绘制所有单条曲线（透明）
for i, d in enumerate(all_data):
    ax_phi.plot(np.arange(len(d['phi'])), d['phi'], color=color_individual, alpha=0.15, linewidth=1.0)

# 绘制均值曲线（加粗）
ax_phi.plot(frames, phi_mean, color=color_mean, linewidth=2.5, label='Mean')

# 绘制95%置信区间
ax_phi.fill_between(frames, phi_mean - 1.96 * phi_std, phi_mean + 1.96 * phi_std, 
                    color=color_ci, alpha=0.2, label='95% CI')

ax_phi.set_xlim(0, max_len)
ax_phi.set_ylim(-10, 10)
ax_phi.legend(loc='upper right', fontsize=10, framealpha=0.95)

# === 子图3: Pitch Angle (theta) ===
ax_theta = fig.add_subplot(gs[2])
ax_theta.set_facecolor(subplot_bg)
ax_theta.set_xlabel("Time Step", fontsize=11, fontweight='bold')
ax_theta.set_ylabel("Pitch Angle, $\\theta$ (deg)", fontsize=11, fontweight='bold')
ax_theta.set_title(f"(c) Vehicle Pitch Angle - {len(all_data)} Trajectories", fontsize=12, fontweight='bold', pad=10)
ax_theta.grid(True, alpha=0.35, linewidth=0.8, which='major')
ax_theta.minorticks_on()
ax_theta.grid(True, alpha=0.15, linewidth=0.5, which='minor')

# 绘制所有单条曲线（透明）
for i, d in enumerate(all_data):
    ax_theta.plot(np.arange(len(d['theta'])), d['theta'], color=color_individual, alpha=0.15, linewidth=1.0)

# 绘制均值曲线（加粗）
ax_theta.plot(frames, theta_mean, color=color_mean, linewidth=2.5, label='Mean')

# 绘制95%置信区间
ax_theta.fill_between(frames, theta_mean - 1.96 * theta_std, theta_mean + 1.96 * theta_std, 
                      color=color_ci, alpha=0.2, label='95% CI')

ax_theta.set_xlim(0, max_len)
ax_theta.set_ylim(-10, 10)
ax_theta.legend(loc='upper right', fontsize=10, framealpha=0.95)

# === 子图4: Sideslip Angle (beta) ===
ax_beta = fig.add_subplot(gs[3])
ax_beta.set_facecolor(subplot_bg)
ax_beta.set_xlabel("Time Step", fontsize=11, fontweight='bold')
ax_beta.set_ylabel("Sideslip Angle, $\\beta$ (deg)", fontsize=11, fontweight='bold')
ax_beta.set_title(f"(d) Sideslip Angle - {len(all_data)} Trajectories", fontsize=12, fontweight='bold', pad=10)
ax_beta.grid(True, alpha=0.35, linewidth=0.8, which='major')
ax_beta.minorticks_on()
ax_beta.grid(True, alpha=0.15, linewidth=0.5, which='minor')

# 绘制所有单条曲线（透明）
for i, d in enumerate(all_data):
    ax_beta.plot(np.arange(len(d['beta'])), d['beta'], color=color_individual, alpha=0.15, linewidth=1.0)

# 绘制均值曲线（加粗）
ax_beta.plot(frames, beta_mean, color=color_mean, linewidth=2.5, label='Mean')

# 绘制95%置信区间
ax_beta.fill_between(frames, beta_mean - 1.96 * beta_std, beta_mean + 1.96 * beta_std, 
                     color=color_ci, alpha=0.2, label='95% CI')

ax_beta.set_xlim(0, max_len)
ax_beta.set_ylim(-10, 10)
ax_beta.legend(loc='upper right', fontsize=10, framealpha=0.95)

# === 4. 调整布局并显示 ===
plt.tight_layout(pad=1.5)

# 添加整体标题
fig.suptitle(f'Vehicle Dynamics Analysis - {len(all_data)} Trajectories from {os.path.basename(model_dir)}', 
             fontsize=14, fontweight='bold', y=1.02)

# 保存图片
output_path = os.path.join(model_dir, "trajectory_analysis.png")
plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
print(f"\nSaved figure to: {output_path}")

plt.show()

print("="*60)
print("Visualization complete!")
print(f"Analyzed {len(all_data)} trajectories from {model_dir}")
print("="*60)
